from typing import Dict, Any, List
from langchain.schema import Document
import yaml


class RAGNodes:
    """RAG图的节点定义"""

    def __init__(self, llm, vector_store, config_path: str = "config.yaml"):
        self.llm = llm
        self.vector_store = vector_store

        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        self.rag_config = config['rag']

    def retrieve_documents(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """检索相关文档"""
        query = state.get("query", "")

        # 执行向量检索
        docs_with_scores = self.vector_store.similarity_search_with_score(
            query, k=self.rag_config['top_k']
        )

        # 过滤低分文档
        filtered_docs = [
            doc for doc, score in docs_with_scores
            if score >= self.rag_config['similarity_threshold']
        ]

        state["retrieved_documents"] = filtered_docs
        state["retrieval_scores"] = [
            score for _, score in docs_with_scores
            if score >= self.rag_config['similarity_threshold']
        ]

        return state

    def generate_answer(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """基于检索文档生成答案"""
        query = state.get("query", "")
        docs = state.get("retrieved_documents", [])

        if not docs:
            state["answer"] = "抱歉，没有找到相关信息来回答您的问题。"
            return state

        # 构建上下文
        context = "\n\n".join([doc.page_content for doc in docs])

        # 构建提示词
        prompt = f"""基于以下上下文信息回答问题：

上下文：
{context}

问题：{query}

请基于上下文信息提供准确、详细的回答。如果上下文中没有足够信息回答问题，请说明。

回答："""

        # 生成答案
        # answer = self.llm(prompt)
        # state["answer"] = answer
        # 使用流式输出
        def generate_stream():
            for chunk in self.llm(prompt):
                if chunk and chunk.get("data"):
                    yield chunk["data"]

        return {
            **state,
            "answer": generate_stream(),  # 返回生成器而不是字符串
            "retrieved_documents": docs,
            "retrieval_scores": state.get("retrieval_scores", [])
        }

    def direct_answer(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """直接回答（不需要检索的情况）"""
        state["answer"] = "请提供一个具体的问题，我会为您查找相关信息。"
        return state